package v4.text;

import java.io.UnsupportedEncodingException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AesKit {

	public static final int AES_KEYSIZE_128 = 128;
	public static final int AES_KEYSIZE_192 = 192;
	public static final int AES_KEYSIZE_256 = 256;
	public static final String ECB_PKCS5Padding = "AES/ECB/PKCS5Padding";
	public static final String CBC_PKCS5Padding = "AES/CBC/PKCS5Padding";
	
	public static final String AES = "AES";
	public static final String AES_CIPHER_ALGORITHM = ECB_PKCS5Padding;
	private static final String SHA1PRNG = "SHA1PRNG";
	
	public AesKit() {
		
	}
	public AesKit(int aes_keysize, String aes_padding) {
		this.key_size = aes_keysize;
		this.cipher_algorithm = aes_padding;
	}
	
	private int key_size;
	private String cipher_algorithm;
	
	
	public String aes_encrypt(String content, String password) {
		try {
			byte[] byteContent = content.getBytes(Config.UTF8);
			
			Cipher cipher = Cipher.getInstance(cipher_algorithm);
			if(cipher_algorithm.contains("ECB")) {
				cipher.init(Cipher.ENCRYPT_MODE, getSecretKey(password));
			}else if(cipher_algorithm.contains("CBC")) {
				int len = cipher.getBlockSize();
				byte[] ivByte = new byte[len];
				byte[] kByte = new byte[len];
				byte[] pwByte = password.getBytes(Config.UTF8);
				if(pwByte.length>len) {
					System.arraycopy(pwByte, 0, ivByte, 0, ivByte.length);
					int klen = kByte.length < pwByte.length - ivByte.length ? kByte.length : pwByte.length - ivByte.length;
					System.arraycopy(pwByte, ivByte.length, kByte, 0, klen);
				}else if(pwByte.length<=len) {
					System.arraycopy(pwByte, 0, ivByte, 0, pwByte.length);
					System.arraycopy(pwByte, 0, kByte, 0, pwByte.length);
				}
				IvParameterSpec iv = new IvParameterSpec(ivByte);
				cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(kByte, AES), iv);
			}

			byte[] result = cipher.doFinal(byteContent);
			return Base64.getEncoder().encodeToString(result);
		} catch (Exception e) {
			e.printStackTrace();
			throw new RuntimeException(e.getMessage());
		}

	}

	public String aes_decrypt(String content, String password) {
		try {
			byte[] data = Base64.getDecoder().decode(content);

			Cipher cipher = Cipher.getInstance(cipher_algorithm);
			if(cipher_algorithm.contains("ECB")) {
				cipher.init(Cipher.DECRYPT_MODE, getSecretKey(password));
			}else if(cipher_algorithm.contains("CBC")) {
				int len = cipher.getBlockSize();
				byte[] ivByte = new byte[len];
				byte[] kByte = new byte[len];
				byte[] pwByte = password.getBytes(Config.UTF8);
				if(pwByte.length>len) {
					System.arraycopy(pwByte, 0, ivByte, 0, ivByte.length);
					int klen = kByte.length < pwByte.length - ivByte.length ? kByte.length : pwByte.length - ivByte.length;
					System.arraycopy(pwByte, ivByte.length, kByte, 0, klen);
				}else if(pwByte.length<=len) {
					System.arraycopy(pwByte, 0, ivByte, 0, pwByte.length);
					System.arraycopy(pwByte, 0, kByte, 0, pwByte.length);
				}
				IvParameterSpec iv = new IvParameterSpec(ivByte);
				cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kByte, AES), iv);
			}

			byte[] result = cipher.doFinal(data);
			return new String(result, Config.UTF8);
		} catch (Exception e) {
			e.printStackTrace();
			throw new RuntimeException(e.getMessage());
		}

	}

	private SecretKeySpec getSecretKey(final String password) {
		try {
			/*KeyGenerator kg = KeyGenerator.getInstance(AES);
			kg.init(key_size, new SecureRandom(password.getBytes(Config.UTF8)));
			SecretKey secretKey = kg.generateKey();
			return new SecretKeySpec(secretKey.getEncoded(), AES);*/
			
			/*KeyGenerator kg = KeyGenerator.getInstance(AES);
			SecureRandom secureRandom = new SecureRandom();
	        secureRandom.setSeed(password.getBytes(Config.UTF8));
	        kg.init(key_size, secureRandom);
	        SecretKey secretKey = kg.generateKey();
	        return new SecretKeySpec(secretKey.getEncoded(), AES);*/
			
			KeyGenerator kg = KeyGenerator.getInstance(AES);
			SecureRandom sr = SecureRandom.getInstance(SHA1PRNG);
			sr.setSeed(password.getBytes(Config.UTF8));
			kg.init(key_size, sr);
			SecretKey secretKey = kg.generateKey();
			return new SecretKeySpec(secretKey.getEncoded(), AES);
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
			throw new RuntimeException(e.getMessage());
		} catch (UnsupportedEncodingException e) {
			e.printStackTrace();
			throw new RuntimeException(e.getMessage());
		}

	}
	
}
